# multi_drama_generator.py
import re
import os
import json
import time
import logging
from typing import Dict, List, Tuple, Optional, Union
from pathlib import Path
import numpy as np

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class Emotion:
    """情绪枚举类"""
    NEUTRAL = "neutral"  # 平静
    HAPPY = "happy"  # 开心
    SAD = "sad"  # 悲伤
    ANGRY = "angry"  # 愤怒
    EXCITED = "excited"  # 兴奋
    CALM = "calm"  # 冷静
    SURPRISED = "surprised"  # 惊讶


class DialogueLine:
    """对话行数据结构"""

    def __init__(self, character: str, text: str, emotion: str = Emotion.NEUTRAL,
                 speed: float = 1.0, volume: float = 0.8, line_number: int = 0, raw_text: str = ""):
        self.character = character  # 角色名
        self.text = text  # 对话文本
        self.emotion = emotion  # 情绪
        self.speed = speed  # 语速 (0.5-2.0)
        self.volume = volume  # 音量 (0.1-1.0)
        self.line_number = line_number  # 行号
        self.raw_text = raw_text  # 原始文本

    def __str__(self):
        return f"DialogueLine(character='{self.character}', text='{self.text}', emotion='{self.emotion}')"


class ScriptParser:
    """剧本解析器"""

    def __init__(self):
        # 角色名匹配模式 - 支持中英文
        self.character_pattern = r'^([^:：\n]+)[：:](.+)$'

        # 情绪标签模式 - 支持多种格式
        self.emotion_patterns = [
            r'\[([^[\]]+)\]',  # [开心]
            r'（([^（）]+)）',  # （开心）
            r'\(([^()]+)\)',  # (开心)
            r'<([^<>]+)>',  # <开心>
        ]

        # 语速控制模式
        self.speed_pattern = r'@速度[：:]?([\d.]+)'

        # 音量控制模式
        self.volume_pattern = r'@音量[：:]?([\d.]+)'

        # 情绪映射字典
        self.emotion_mapping = {
            '开心': Emotion.HAPPY, '高兴': Emotion.HAPPY, '兴奋': Emotion.EXCITED,
            '悲伤': Emotion.SAD, '难过': Emotion.SAD, '伤心': Emotion.SAD,
            '愤怒': Emotion.ANGRY, '生气': Emotion.ANGRY, '愤慨': Emotion.ANGRY,
            '平静': Emotion.NEUTRAL, '正常': Emotion.NEUTRAL, '普通': Emotion.NEUTRAL,
            '冷静': Emotion.CALM, '沉着': Emotion.CALM,
            '惊讶': Emotion.SURPRISED, '吃惊': Emotion.SURPRISED, '震惊': Emotion.SURPRISED,
            'happy': Emotion.HAPPY, 'excited': Emotion.EXCITED,
            'sad': Emotion.SAD, 'angry': Emotion.ANGRY,
            'neutral': Emotion.NEUTRAL, 'calm': Emotion.CALM,
            'surprised': Emotion.SURPRISED
        }

    def parse_script(self, script_content: str) -> List[DialogueLine]:
        """解析剧本内容"""
        lines = script_content.strip().split('\n')
        dialogue_lines = []

        for line_num, line in enumerate(lines, 1):
            line = line.strip()

            # 跳过空行和注释行
            if not line or line.startswith('#') or line.startswith('//'):
                continue

            # 解析对话行
            dialogue_line = self._parse_dialogue_line(line, line_num)
            if dialogue_line:
                dialogue_lines.append(dialogue_line)

        return dialogue_lines

    def _parse_dialogue_line(self, line: str, line_num: int) -> Optional[DialogueLine]:
        """解析单行对话"""
        # 匹配角色名和对话内容
        match = re.match(self.character_pattern, line.strip())
        if not match:
            logger.warning(f"第{line_num}行格式不正确，跳过: {line}")
            return None

        character = match.group(1).strip()
        content = match.group(2).strip()

        # 解析情绪标签
        emotion, cleaned_content = self._extract_emotion(content)

        # 解析语速控制
        speed, cleaned_content = self._extract_speed(cleaned_content)

        # 解析音量控制
        volume, cleaned_content = self._extract_volume(cleaned_content)

        return DialogueLine(
            character=character,
            text=cleaned_content,
            emotion=emotion,
            speed=speed,
            volume=volume,
            line_number=line_num,
            raw_text=line
        )

    def _extract_emotion(self, text: str) -> Tuple[str, str]:
        """提取情绪标签"""
        for pattern in self.emotion_patterns:
            match = re.search(pattern, text)
            if match:
                emotion_text = match.group(1).strip()
                # 移除情绪标签
                cleaned_text = re.sub(pattern, '', text).strip()

                # 映射情绪
                emotion = self.emotion_mapping.get(emotion_text, Emotion.NEUTRAL)
                return emotion, cleaned_text

        return Emotion.NEUTRAL, text

    def _extract_speed(self, text: str) -> Tuple[float, str]:
        """提取语速控制"""
        match = re.search(self.speed_pattern, text)
        if match:
            try:
                speed = float(match.group(1))
                # 限制语速范围
                speed = max(0.5, min(2.0, speed))
                # 移除语速标签
                cleaned_text = re.sub(self.speed_pattern, '', text).strip()
                return speed, cleaned_text
            except ValueError:
                pass

        return 1.0, text  # 默认语速

    def _extract_volume(self, text: str) -> Tuple[float, str]:
        """提取音量控制"""
        match = re.search(self.volume_pattern, text)
        if match:
            try:
                volume = float(match.group(1))
                # 限制音量范围
                volume = max(0.1, min(1.0, volume))
                # 移除音量标签
                cleaned_text = re.sub(self.volume_pattern, '', text).strip()
                return volume, cleaned_text
            except ValueError:
                pass

        return 0.8, text  # 默认音量

    def get_characters(self, dialogue_lines: List[DialogueLine]) -> List[str]:
        """获取剧本中的所有角色"""
        characters = set()
        for line in dialogue_lines:
            characters.add(line.character)
        return sorted(list(characters))

    def validate_script(self, dialogue_lines: List[DialogueLine]) -> Dict[str, any]:
        """验证剧本格式和内容"""
        validation_result = {
            'is_valid': True,
            'total_lines': len(dialogue_lines),
            'characters': self.get_characters(dialogue_lines),
            'character_count': len(self.get_characters(dialogue_lines)),
            'warnings': [],
            'errors': []
        }

        # 检查是否有对话行
        if not dialogue_lines:
            validation_result['is_valid'] = False
            validation_result['errors'].append("剧本中没有找到有效的对话行")
            return validation_result

        # 检查角色数量
        if validation_result['character_count'] > 10:
            validation_result['warnings'].append(
                f"角色数量较多({validation_result['character_count']}个)，可能影响处理速度")

        # 检查角色名长度
        for char in validation_result['characters']:
            if len(char) > 20:
                validation_result['warnings'].append(f"角色名'{char}'过长，建议简化")

        # 检查对话内容长度
        long_lines = [line for line in dialogue_lines if len(line.text) > 200]
        if long_lines:
            validation_result['warnings'].append(f"发现{len(long_lines)}行对话过长，可能影响语音合成效果")

        return validation_result